import os
import glob
import argparse
import gmalg

def verify_keys (prikey_file, pubkey_file):
    sm2 = gmalg.SM2()
    sk = read_binary_file_(prikey_file)
    pk = read_binary_file_(pubkey_file)
    pk_new  = sm2.generate_pk(sk)
    assert pk == pk_new, "公钥验证失败"

def read_binary_file_(file_path):
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        return binary_data

def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    parser.add_argument('mode', choices=['sign', 'verify'], default='sign', help='运行模式')
    args = parser.parse_args()
    directory = args.directory
    mode = args.mode

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    # 创建输出内容
    merged_content = ""
    
    # 遍历0-9，合并对应的公钥和私钥文件
    for i in range(10):
        pubkey_file = os.path.join(directory, f'sm2_pubkey_{i}')
        prikey_file = os.path.join(directory, f'sm2_prikey_{i}')
        msg_file = os.path.join(directory, f'sm2_msg_{i}')
        sign_file = os.path.join(directory, f'sm2_sign_{i}')
        
        verify_keys (prikey_file, pubkey_file)
        # 读取文件内容
        pubkey_content = read_binary_file(pubkey_file)
        prikey_content = read_binary_file(prikey_file)
        msg_content = read_binary_file(msg_file)
        sign_content = read_binary_file(sign_file)
        
        # 添加到合并内容
        merged_content += f"公钥= \n{pubkey_content}\n"
        merged_content += f"私钥= \n{prikey_content}\n"
        merged_content += f"签名数据= \n{msg_content}\n"
        merged_content += f"签名值= \n{sign_content}\n"
        merged_content += "\n"
    
    # 写入到txt文件
    if mode == 'sign':
        output_file = 'sm2签名无预处理.txt'
    else:
        output_file = 'sm2验签无预处理.txt'

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
    
    print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()

